import json
import sys
from itertools import islice
import subprocess

# 找到倒数第二行
result = subprocess.run(["tail", "-n", "2", sys.argv[1]], 
                        capture_output=True, text=True)
lines = result.stdout.splitlines()
second_last_line = lines[0] if len(lines) >= 2 else None
base_ts = int(second_last_line.split(":")[1].strip())  # 提取并转为整数

with open(sys.argv[1], "r") as f:
    torch_trace_lines = f.readlines()
    insert_pos = len(torch_trace_lines) - 5  # 倒数5行
    torch_trace_lines[insert_pos - 1] = "}"     # 实际上就是要去掉这一行最后的换行符

llm_trace_strs = ""
with open(sys.argv[2], "r") as f:
    while True:
        chunk = list(islice(f, 7))  # 读取 7 行
        if not chunk:  # 如果 chunk 为空，说明文件读完
            break
        trace_obj = {}
        trace_obj["ph"] = "X"
        trace_obj["cat"] = chunk[0].strip()
        trace_obj["name"] = chunk[1].strip()
        trace_obj["pid"] = int(chunk[2].strip())
        trace_obj["tid"] = int(chunk[3].strip())
        trace_obj["ts"] = (int(chunk[4].strip()) - base_ts) / 1000.0
        trace_obj["dur"] = int(chunk[5].strip()) / 1000.0    # 保留3位小数
        trace_obj["args"] = {"args":chunk[6].strip()}
        llm_trace_strs += ",\n" + json.dumps(trace_obj, indent=4)

torch_trace_lines.insert(insert_pos, llm_trace_strs + "\n")
with open("merged_trace.json", "w") as f:
    f.writelines(torch_trace_lines)

